import json
import tushare as ts
from datetime import datetime
import pandas as pd

class StockDataExporter:
    def __init__(self):
        # Tushare初始化
        self.pro = ts.pro_api('20241126212238-f4c9b74f-0059-4a93-9d1f-53dae167179a')
        self.pro._DataApi__http_url = 'http://tsapi.majors.ltd:7000'
    
    def fetch_stock_data(self):
        """获取股票数据"""
        try:
            print(f"开始获取股票数据，时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
            df = self.pro.stock_basic(
                fields=[
                    "ts_code", "symbol", "name", "area", "industry",
                    "cnspell", "market", "list_date", "act_name",
                    "act_ent_type", "fullname", "enname", "exchange",
                    "curr_type", "list_status", "delist_date", "is_hs"
                ]
            )
            
            if df is not None and not df.empty:
                print(f"成功获取到 {len(df)} 条股票数据")
                print("\n数据样例（前3条）：")
                print(df[['ts_code', 'name', 'market']].head(3))
            else:
                print("警告：获取到的数据为空")
                
            return df
        except Exception as e:
            print(f"获取数据失败: {str(e)}")
            return None

    def process_for_unicloud(self, df):
        """处理数据为uniCloud格式"""
        records = []
        for _, row in df.iterrows():
            # 创建基础记录
            record = {
                "ts_code": row['ts_code'],          # 保留原始 ts_code
                "stock_code": row['ts_code'],       # 新增 stock_code 字段
                "symbol": row['symbol'],
                "name": row['name'],
                "area": row['area'] if pd.notna(row['area']) else "",
                "industry": row['industry'] if pd.notna(row['industry']) else "",
                "market": row['market'],
                "exchange": row['exchange'],
                "list_status": row['list_status'],
                "is_hs": row['is_hs'] if pd.notna(row['is_hs']) else "N",
                "create_time": {"$date": datetime.now().strftime("%Y-%m-%dT%H:%M:%SZ")},
                "update_time": {"$date": datetime.now().strftime("%Y-%m-%dT%H:%M:%SZ")}
            }
            
            # 处理可选字段
            optional_fields = [
                'fullname', 'enname', 'cnspell', 'curr_type',
                'act_name', 'act_ent_type'
            ]
            
            for field in optional_fields:
                if pd.notna(row[field]):
                    record[field] = row[field]
            
            # 处理上市日期，确保正确的格式
            if pd.notna(row['list_date']):
                date_str = str(row['list_date'])
                formatted_date = f"{date_str[:4]}-{date_str[4:6]}-{date_str[6:]}T00:00:00Z"
                record['list_date'] = {"$date": formatted_date}
            
            if pd.notna(row['delist_date']):
                date_str = str(row['delist_date'])
                formatted_date = f"{date_str[:4]}-{date_str[4:6]}-{date_str[6:]}T00:00:00Z"
                record['delist_date'] = {"$date": formatted_date}
            
            records.append(record)
        
        return records

    def export_to_jsonl(self, records, filename='stock_company_basic_data.json'):
        """导出为JSONL格式"""
        try:
            record_count = len(records)
            print(f"准备导出 {record_count} 条记录...")
            
            with open(filename, 'w', encoding='utf-8') as f:
                for i, record in enumerate(records, 1):
                    f.write(json.dumps(record, ensure_ascii=False) + '\n')
                    if i % 1000 == 0:
                        print(f"已处理 {i}/{record_count} 条记录...")
                        
            print(f"成功导出数据到 {filename}")
            # 验证文件大小
            import os
            file_size = os.path.getsize(filename) / 1024 / 1024
            print(f"导出文件大小：{file_size:.2f}MB")
            
        except Exception as e:
            print(f"导出数据失败: {str(e)}")
        
        
    def run(self):
        """运行导出流程"""
        print("开始获取股票数据...")
        df = self.fetch_stock_data()
        
        if df is None:
            print("未能获取到股票数据，请检查 tushare token 是否有效")
            return
            
        if df.empty:
            print("获取到的数据集为空")
            return
        
        print(f"获取到 {len(df)} 条股票数据")
        print("处理数据格式...")
        records = self.process_for_unicloud(df)
        
        if not records:
            print("处理后的记录为空")
            return
        
        print("导出数据...")
        self.export_to_jsonl(records)
        
        print(f"共导出 {len(records)} 条记录")
        print("导出完成")

if __name__ == "__main__":
    exporter = StockDataExporter()
    exporter.run()